/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com> 
 *
 */
#ifndef SRC_MULTIPLEXER_H_
#define SRC_MULTIPLEXER_H_

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "IPAddress.h"
#include "Packet.h"
#if defined(IS_FREEBSD)
#include <sys/socket.h>
#endif
#include <iostream>
#include <memory>
#include <functional>
#include <unordered_map>

namespace aiengine {

class Protocol;
typedef std::shared_ptr<Protocol> ProtocolPtr;

#define NO_PROTOCOL_SELECTED 0xFFFF

class Multiplexer;
typedef std::shared_ptr<Multiplexer> MultiplexerPtr;
typedef std::weak_ptr<Multiplexer> MultiplexerPtrWeak;

class Multiplexer {
public:
    	explicit Multiplexer() = default; 
    	virtual ~Multiplexer() {}

	void addUpMultiplexer(MultiplexerPtrWeak mux);
	void addDownMultiplexer(MultiplexerPtrWeak mux);

	MultiplexerPtrWeak getDownMultiplexer() const;
	MultiplexerPtrWeak getUpMultiplexer(int key) const;

	void forward(Packet& packet);

        void statistics(std::basic_ostream<char> &out);
        void statistics() { statistics(std::cout); }

	int getNumberUpMultiplexers() const { return (int)muxUpMap_.size(); }

	uint16_t getProtocolIdentifier() const { return protocol_id_; }
	uint16_t getNextProtocolIdentifier() const { return next_protocol_id_; }
	void setNextProtocolIdentifier(uint16_t protocol_id) { next_protocol_id_ = protocol_id; }
	void setProtocol(const ProtocolPtr& proto);
	void setProtocol(const ProtocolPtr& proto, uint16_t protocol_id);
	ProtocolPtr getProtocol() const { return proto_; }

	void setHeaderSize(int size) { header_size_ = size; }
	int getHeaderSize() const { return header_size_; }

	void setPacketInfo(const uint8_t *packet, int length, int prev_header_size, PacketAnomalyType pa,time_t packet_time);
	void setPacket(Packet *packet);

	uint64_t getTotalForwardPackets() const { return total_forward_packets_; }
	uint64_t getTotalFailPackets() const { return total_fail_packets_; }
	uint64_t getTotalReceivedPackets() const { return total_received_packets_; }

	Packet *getCurrentPacket() { return &packet_;}

	bool accept(Packet &packet) const { return check_func_(packet); }

	IPAddress address;
	// This is realy uggly puagggggg
	uint16_t total_length = 0;
private:
	ProtocolPtr proto_ = nullptr;
	Packet packet_ {};
	uint64_t total_received_packets_ = 0;
	uint64_t total_forward_packets_ = 0;
	uint64_t total_fail_packets_ = 0;
	MultiplexerPtrWeak muxDown_ = MultiplexerPtrWeak();
	int header_size_ = 0;
	uint16_t protocol_id_ = NO_PROTOCOL_SELECTED; // the protocol analiyzer owned by the multiplexer
	uint16_t next_protocol_id_ = NO_PROTOCOL_SELECTED; // the next protocol to check by the multiplexer
    	typedef std::unordered_map<int,MultiplexerPtrWeak> MuxMap;
	MuxMap muxUpMap_ {};
	std::function <bool (Packet&)> check_func_ = [&] (Packet&) { return true; };
	std::function <bool (Packet&)> packet_func_ = [&] (Packet&) { return true; };
};

} // namespace aiengine

#endif  // SRC_MULTIPLEXER_H_
